package splitter;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * 基于抽象语法树的代码分割器实现
 */
public class AstCodeSplitter implements Splitter {
    private int chunkSize = 2500;
    private int chunkOverlap = 300;
    private LangChainCodeSplitter langchainFallback;
    
    // 可分割的节点类型
    private static final Map<String, List<String>> SPLITTABLE_NODE_TYPES = new HashMap<>();
    
    static {
        // 初始化可分割的节点类型
        SPLITTABLE_NODE_TYPES.put("javascript", Arrays.asList(
            "function_declaration", "arrow_function", "class_declaration", "method_definition", "export_statement"
        ));
        SPLITTABLE_NODE_TYPES.put("typescript", Arrays.asList(
            "function_declaration", "arrow_function", "class_declaration", "method_definition", "export_statement",
            "interface_declaration", "type_alias_declaration"
        ));
        SPLITTABLE_NODE_TYPES.put("python", Arrays.asList(
            "function_definition", "class_definition", "decorated_definition", "async_function_definition"
        ));
        SPLITTABLE_NODE_TYPES.put("java", Arrays.asList(
            "method_declaration", "class_declaration", "interface_declaration", "constructor_declaration"
        ));
        SPLITTABLE_NODE_TYPES.put("cpp", Arrays.asList(
            "function_definition", "class_specifier", "namespace_definition", "declaration"
        ));
        SPLITTABLE_NODE_TYPES.put("go", Arrays.asList(
            "function_declaration", "method_declaration", "type_declaration", "var_declaration", "const_declaration"
        ));
        SPLITTABLE_NODE_TYPES.put("rust", Arrays.asList(
            "function_item", "impl_item", "struct_item", "enum_item", "trait_item", "mod_item"
        ));
    }

    /**
     * 创建一个新的 AST 代码分割器
     */
    public AstCodeSplitter() {
        this.langchainFallback = new LangChainCodeSplitter(this.chunkSize, this.chunkOverlap);
    }

    /**
     * 创建一个新的 AST 代码分割器
     * 
     * @param chunkSize 代码块大小
     * @param chunkOverlap 代码块重叠大小
     */
    public AstCodeSplitter(int chunkSize, int chunkOverlap) {
        if (chunkSize > 0) this.chunkSize = chunkSize;
        if (chunkOverlap >= 0) this.chunkOverlap = chunkOverlap;
        this.langchainFallback = new LangChainCodeSplitter(this.chunkSize, this.chunkOverlap);
    }

    @Override
    public CompletableFuture<List<CodeChunk>> split(String code, String language, String filePath) throws Exception {
        return CompletableFuture.supplyAsync(() -> {
            try {
                // 检查语言是否受 AST 分割器支持
                if (!isLanguageSupported(language)) {
                    System.out.println("📝 Language " + language + " not supported by AST, using LangChain splitter for: " + 
                        (filePath != null ? filePath : "unknown"));
                    return langchainFallback.split(code, language, filePath).join();
                }
                
                System.out.println("🌳 Using AST splitter for " + language + " file: " + 
                    (filePath != null ? filePath : "unknown"));
                
                // 在 Java 实现中，我们需要使用 Java 解析器库来解析 AST
                // 这里我们使用简化的基于正则表达式的方法来模拟 AST 分割
                List<CodeChunk> chunks = extractChunks(code, language, filePath);
                
                // 如果块太大，进一步分割
                List<CodeChunk> refinedChunks = refineChunks(chunks, code);
                
                return refinedChunks;
            } catch (Exception e) {
                System.err.println("⚠️ AST splitter failed for " + language + ", falling back to LangChain: " + e.getMessage());
                e.printStackTrace();
                return langchainFallback.split(code, language, filePath).join();
            }
        });
    }

    @Override
    public void setChunkSize(int chunkSize) {
        this.chunkSize = chunkSize;
        this.langchainFallback.setChunkSize(chunkSize);
    }

    @Override
    public void setChunkOverlap(int chunkOverlap) {
        this.chunkOverlap = chunkOverlap;
        this.langchainFallback.setChunkOverlap(chunkOverlap);
    }

    /**
     * 提取代码块
     * 
     * @param code 代码内容
     * @param language 语言
     * @param filePath 文件路径
     * @return 代码块列表
     */
    private List<CodeChunk> extractChunks(String code, String language, String filePath) {
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        // 根据不同语言使用不同的分割策略
        switch (language.toLowerCase()) {
            case "javascript":
            case "typescript":
            case "js":
            case "ts":
                chunks = extractJavaScriptChunks(code, language, filePath);
                break;
            case "python":
            case "py":
                chunks = extractPythonChunks(code, language, filePath);
                break;
            case "java":
                chunks = extractJavaChunks(code, language, filePath);
                break;
            case "cpp":
            case "c++":
            case "c":
                chunks = extractCppChunks(code, language, filePath);
                break;
            case "go":
                chunks = extractGoChunks(code, language, filePath);
                break;
            case "rust":
            case "rs":
                chunks = extractRustChunks(code, language, filePath);
                break;
            default:
                // 对于不支持的语言，使用 LangChain 分割器
                return langchainFallback.split(code, language, filePath).join();
        }
        
        // 如果没有找到有意义的块，创建一个包含整个代码的块
        if (chunks.isEmpty()) {
            chunks.add(new CodeChunk(
                code,
                1,
                lines.length,
                language,
                filePath
            ));
        }
        
        return chunks;
    }

    /**
     * 提取 JavaScript/TypeScript 代码块
     */
    private List<CodeChunk> extractJavaScriptChunks(String code, String language, String filePath) {
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        // 使用正则表达式模拟 AST 解析
        // 在实际实现中，应该使用专门的 JavaScript 解析器
        
        int lineNumber = 1;
        int startLine = 1;
        StringBuilder currentChunk = new StringBuilder();
        
        for (String line : lines) {
            // 检查是否是函数、类或方法定义的开始
            if (line.matches(".*\\b(function|class|interface|type)\\s+\\w+.*") ||
                line.matches(".*\\bconst\\s+\\w+\\s*=\\s*\\(.*") ||
                line.matches(".*\\bexport\\s+(function|class|interface|type|const).*")) {
                
                // 如果已经有内容，保存当前块
                if (currentChunk.length() > 0) {
                    chunks.add(new CodeChunk(
                        currentChunk.toString(),
                        startLine,
                        lineNumber - 1,
                        language,
                        filePath
                    ));
                    
                    currentChunk = new StringBuilder();
                    startLine = lineNumber;
                }
            }
            
            currentChunk.append(line).append("\n");
            lineNumber++;
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                lineNumber - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }
    
    /**
     * 提取 Python 代码块
     */
    private List<CodeChunk> extractPythonChunks(String code, String language, String filePath) {
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        int lineNumber = 1;
        int startLine = 1;
        StringBuilder currentChunk = new StringBuilder();
        
        for (String line : lines) {
            // 检查是否是函数或类定义的开始
            if (line.matches("^\\s*(def|class)\\s+\\w+.*") ||
                line.matches("^\\s*@.*\\s*$")) {
                
                // 如果已经有内容，保存当前块
                if (currentChunk.length() > 0) {
                    chunks.add(new CodeChunk(
                        currentChunk.toString(),
                        startLine,
                        lineNumber - 1,
                        language,
                        filePath
                    ));
                    
                    currentChunk = new StringBuilder();
                    startLine = lineNumber;
                }
            }
            
            currentChunk.append(line).append("\n");
            lineNumber++;
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                lineNumber - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }
    
    /**
     * 提取 Java 代码块
     */
    private List<CodeChunk> extractJavaChunks(String code, String language, String filePath) {
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        int lineNumber = 1;
        int startLine = 1;
        StringBuilder currentChunk = new StringBuilder();
        
        for (String line : lines) {
            // 检查是否是方法、类或接口定义的开始
            if (line.matches("^\\s*(public|private|protected|static|\\s)*\\s*\\w+\\s+\\w+\\s*\\([^)]*\\)\\s*\\{?.*") ||
                line.matches("^\\s*(public|private|protected|static|\\s)*\\s*(class|interface|enum)\\s+\\w+.*")) {
                
                // 如果已经有内容，保存当前块
                if (currentChunk.length() > 0) {
                    chunks.add(new CodeChunk(
                        currentChunk.toString(),
                        startLine,
                        lineNumber - 1,
                        language,
                        filePath
                    ));
                    
                    currentChunk = new StringBuilder();
                    startLine = lineNumber;
                }
            }
            
            currentChunk.append(line).append("\n");
            lineNumber++;
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                lineNumber - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }
    
    /**
     * 提取 C/C++ 代码块
     */
    private List<CodeChunk> extractCppChunks(String code, String language, String filePath) {
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        int lineNumber = 1;
        int startLine = 1;
        StringBuilder currentChunk = new StringBuilder();
        
        for (String line : lines) {
            // 检查是否是函数、类或命名空间定义的开始
            if (line.matches("^\\s*\\w+\\s+\\w+\\s*\\([^)]*\\)\\s*\\{?.*") ||
                line.matches("^\\s*(class|struct|namespace)\\s+\\w+.*")) {
                
                // 如果已经有内容，保存当前块
                if (currentChunk.length() > 0) {
                    chunks.add(new CodeChunk(
                        currentChunk.toString(),
                        startLine,
                        lineNumber - 1,
                        language,
                        filePath
                    ));
                    
                    currentChunk = new StringBuilder();
                    startLine = lineNumber;
                }
            }
            
            currentChunk.append(line).append("\n");
            lineNumber++;
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                lineNumber - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }
    
    /**
     * 提取 Go 代码块
     */
    private List<CodeChunk> extractGoChunks(String code, String language, String filePath) {
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        int lineNumber = 1;
        int startLine = 1;
        StringBuilder currentChunk = new StringBuilder();
        
        for (String line : lines) {
            // 检查是否是函数、方法或类型定义的开始
            if (line.matches("^\\s*func\\s+\\w+.*") ||
                line.matches("^\\s*func\\s*\\([^)]*\\)\\s*\\w+.*") ||
                line.matches("^\\s*type\\s+\\w+.*")) {
                
                // 如果已经有内容，保存当前块
                if (currentChunk.length() > 0) {
                    chunks.add(new CodeChunk(
                        currentChunk.toString(),
                        startLine,
                        lineNumber - 1,
                        language,
                        filePath
                    ));
                    
                    currentChunk = new StringBuilder();
                    startLine = lineNumber;
                }
            }
            
            currentChunk.append(line).append("\n");
            lineNumber++;
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                lineNumber - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }
    
    /**
     * 提取 Rust 代码块
     */
    private List<CodeChunk> extractRustChunks(String code, String language, String filePath) {
        List<CodeChunk> chunks = new ArrayList<>();
        String[] lines = code.split("\n");
        
        int lineNumber = 1;
        int startLine = 1;
        StringBuilder currentChunk = new StringBuilder();
        
        for (String line : lines) {
            // 检查是否是函数、结构体、枚举、特性或模块定义的开始
            if (line.matches("^\\s*fn\\s+\\w+.*") ||
                line.matches("^\\s*struct\\s+\\w+.*") ||
                line.matches("^\\s*enum\\s+\\w+.*") ||
                line.matches("^\\s*trait\\s+\\w+.*") ||
                line.matches("^\\s*impl\\s+.*") ||
                line.matches("^\\s*mod\\s+\\w+.*")) {
                
                // 如果已经有内容，保存当前块
                if (currentChunk.length() > 0) {
                    chunks.add(new CodeChunk(
                        currentChunk.toString(),
                        startLine,
                        lineNumber - 1,
                        language,
                        filePath
                    ));
                    
                    currentChunk = new StringBuilder();
                    startLine = lineNumber;
                }
            }
            
            currentChunk.append(line).append("\n");
            lineNumber++;
        }
        
        // 添加最后一个块
        if (currentChunk.length() > 0) {
            chunks.add(new CodeChunk(
                currentChunk.toString(),
                startLine,
                lineNumber - 1,
                language,
                filePath
            ));
        }
        
        return chunks;
    }

    /**
     * 进一步处理代码块，分割大块并添加重叠
     * 
     * @param chunks 原始代码块
     * @param originalCode 原始代码
     * @return 处理后的代码块
     */
    private List<CodeChunk> refineChunks(List<CodeChunk> chunks, String originalCode) {
        List<CodeChunk> refinedChunks = new ArrayList<>();
        
        for (CodeChunk chunk : chunks) {
            if (chunk.getContent().length() <= this.chunkSize) {
                refinedChunks.add(chunk);
            } else {
                // 分割大块
                List<CodeChunk> subChunks = splitLargeChunk(chunk, originalCode);
                refinedChunks.addAll(subChunks);
            }
        }
        
        return addOverlap(refinedChunks);
    }

    /**
     * 分割大代码块
     * 
     * @param chunk 大代码块
     * @param originalCode 原始代码
     * @return 分割后的代码块列表
     */
    private List<CodeChunk> splitLargeChunk(CodeChunk chunk, String originalCode) {
        List<CodeChunk> subChunks = new ArrayList<>();
        String[] lines = chunk.getContent().split("\n");
        
        StringBuilder currentChunk = new StringBuilder();
        int currentStartLine = chunk.getMetadata().getStartLine();
        int currentLineCount = 0;
        
        for (int i = 0; i < lines.length; i++) {
            String line = lines[i];
            String lineWithNewline = i == lines.length - 1 ? line : line + "\n";
            
            if (currentChunk.length() + lineWithNewline.length() > this.chunkSize && currentChunk.length() > 0) {
                // 创建子块
                subChunks.add(new CodeChunk(
                    currentChunk.toString().trim(),
                    currentStartLine,
                    currentStartLine + currentLineCount - 1,
                    chunk.getMetadata().getLanguage(),
                    chunk.getMetadata().getFilePath()
                ));
                
                currentChunk = new StringBuilder();
                currentStartLine = chunk.getMetadata().getStartLine() + i;
                currentLineCount = 1;
            } else {
                currentChunk.append(lineWithNewline);
                currentLineCount++;
            }
        }
        
        // 添加最后一个子块
        if (currentChunk.length() > 0) {
            subChunks.add(new CodeChunk(
                currentChunk.toString().trim(),
                currentStartLine,
                currentStartLine + currentLineCount - 1,
                chunk.getMetadata().getLanguage(),
                chunk.getMetadata().getFilePath()
            ));
        }
        
        return subChunks;
    }

    /**
     * 添加代码块重叠
     * 
     * @param chunks 代码块列表
     * @return 添加重叠后的代码块列表
     */
    private List<CodeChunk> addOverlap(List<CodeChunk> chunks) {
        if (chunks.size() <= 1 || this.chunkOverlap <= 0) {
            return chunks;
        }
        
        List<CodeChunk> overlappedChunks = new ArrayList<>();
        
        for (int i = 0; i < chunks.size(); i++) {
            String content = chunks.get(i).getContent();
            CodeChunk.ChunkMetadata metadata = new CodeChunk.ChunkMetadata(
                chunks.get(i).getMetadata().getStartLine(),
                chunks.get(i).getMetadata().getEndLine(),
                chunks.get(i).getMetadata().getLanguage(),
                chunks.get(i).getMetadata().getFilePath()
            );
            
            // 添加前一个块的重叠
            if (i > 0 && this.chunkOverlap > 0) {
                CodeChunk prevChunk = chunks.get(i - 1);
                String overlapText = getLastNChars(prevChunk.getContent(), this.chunkOverlap);
                content = overlapText + "\n" + content;
                metadata.setStartLine(Math.max(1, metadata.getStartLine() - getLineCount(overlapText)));
            }
            
            overlappedChunks.add(new CodeChunk(content, metadata));
        }
        
        return overlappedChunks;
    }

    /**
     * 获取文本的最后 N 个字符
     * 
     * @param text 文本
     * @param n 字符数
     * @return 最后 N 个字符
     */
    private String getLastNChars(String text, int n) {
        if (text.length() <= n) {
            return text;
        }
        return text.substring(text.length() - n);
    }

    /**
     * 获取文本的行数
     * 
     * @param text 文本
     * @return 行数
     */
    private int getLineCount(String text) {
        return text.split("\n").length;
    }

    /**
     * 检查语言是否受 AST 分割器支持
     * 
     * @param language 语言
     * @return 是否支持
     */
    public static boolean isLanguageSupported(String language) {
        String lowerLang = language.toLowerCase();
        return lowerLang.equals("javascript") || lowerLang.equals("js") ||
               lowerLang.equals("typescript") || lowerLang.equals("ts") ||
               lowerLang.equals("python") || lowerLang.equals("py") ||
               lowerLang.equals("java") ||
               lowerLang.equals("cpp") || lowerLang.equals("c++") || lowerLang.equals("c") ||
               lowerLang.equals("go") ||
               lowerLang.equals("rust") || lowerLang.equals("rs");
    }

    /**
     * 获取支持的语言列表
     * 
     * @return 支持的语言列表
     */
    public static List<String> getSupportedLanguages() {
        return Arrays.asList(
            "javascript", "typescript", "python", "java", "cpp", "c++", "c", "go", "rust"
        );
    }
}